#!/usr/bin/env python3
"""
json_to_jsonl.py

Usage:
    python json_to_jsonl.py input.json output.jsonl
    python json_to_jsonl.py input.json         # 输出到 input.jsonl

This script:
 - detects if file is JSON array, dict with list under common keys, single object, or already JSONL
 - writes each JSON object as one line (no pretty indent) using ensure_ascii=False
"""

import json
import argparse
import sys
from pathlib import Path

COMMON_LIST_KEYS = ("data", "items", "results", "questions", "entries")

def try_parse_whole_json(text):
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None

def try_parse_jsonl_lines(text):
    objs = []
    for i, raw in enumerate(text.splitlines()):
        s = raw.strip()
        if not s:
            continue
        try:
            objs.append(json.loads(s))
        except json.JSONDecodeError:
            # failed to parse a line -> treat whole file as not jsonl
            return None
    return objs

def convert_file(in_path: Path, out_path: Path, encoding="utf-8"):
    text = in_path.read_text(encoding=encoding)

    # 1) try parse full JSON
    parsed = try_parse_whole_json(text)
    if parsed is not None:
        # if list -> each element
        if isinstance(parsed, list):
            objs = parsed
        elif isinstance(parsed, dict):
            # try common keys
            objs = None
            for k in COMMON_LIST_KEYS:
                if k in parsed and isinstance(parsed[k], list):
                    objs = parsed[k]
                    break
            if objs is None:
                # fallback: write the dict itself as single jsonl object
                objs = [parsed]
        else:
            # primitive -> wrap
            objs = [parsed]

        with out_path.open("w", encoding=encoding, newline="\n") as fout:
            for obj in objs:
                fout.write(json.dumps(obj, ensure_ascii=False, separators=(",", ":")))
                fout.write("\n")
        return len(objs)

    # 2) maybe file is already jsonl (line-delimited JSON)
    parsed_lines = try_parse_jsonl_lines(text)
    if parsed_lines is not None:
        with out_path.open("w", encoding=encoding, newline="\n") as fout:
            for obj in parsed_lines:
                fout.write(json.dumps(obj, ensure_ascii=False, separators=(",", ":")))
                fout.write("\n")
        return len(parsed_lines)

    # 3) unknown format -> error
    raise ValueError("无法识别输入文件为有效 JSON 或 JSONL 格式。请检查文件。")

def main():
    p = argparse.ArgumentParser(description="Convert JSON -> JSONL (one JSON object per line).")
    p.add_argument("input", help="输入 JSON 或 JSONL 文件路径")
    p.add_argument("output", nargs="?", help="输出 JSONL 路径（默认：输入同名 .jsonl）")
    args = p.parse_args()

    in_path = Path(args.input)
    if not in_path.exists():
        print(f"错误：输入文件不存在: {in_path}", file=sys.stderr)
        sys.exit(2)

    out_path = Path(args.output) if args.output else in_path.with_suffix(".jsonl")
    try:
        count = convert_file(in_path, out_path)
        print(f"已写入 {count} 条记录 到 {out_path}")
    except Exception as e:
        print("转换失败：", e, file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()
